package workpool

import (
	"context"
	"sync"

	"gitee.com/lihaiping1603/utils/log"
)

type Worker interface {
	DoWork(ctx context.Context, args ...interface{})
}

type DoWorkFunc func(ctx context.Context, args ...interface{})

func (f DoWorkFunc) DoWork(ctx context.Context, args ...interface{}) {
	f(ctx, args...)
}

type Work struct {
	worker Worker
	args   []interface{}
}

func NewWork(worker Worker, args ...interface{}) *Work {
	return &Work{worker: worker, args: args}
}

func NewWorkFunc(doworkfn DoWorkFunc, args ...interface{}) *Work {
	return &Work{worker: doworkfn, args: args}
}

type WorkerPool struct {
	numberWorks int
	workerCh    chan *Work
	exit        chan struct{}
	waitGroup   sync.WaitGroup
}

func NewWorkerPool(works int) *WorkerPool {
	return &WorkerPool{
		numberWorks: works,
		workerCh:    make(chan *Work, 1),
		exit:        make(chan struct{}),
	}
}

func (wp *WorkerPool) PutBlocking(work *Work) {
	wp.workerCh <- work
}

func (wp *WorkerPool) PutWithContext(ctx context.Context, work *Work) {
	select {
	case <-ctx.Done():
	case wp.workerCh <- work:
	}
}

func (wp *WorkerPool) Start(ctx context.Context) {
	for i := 0; i < wp.numberWorks; i++ {
		wp.waitGroup.Add(1)
		go func() {
			defer wp.waitGroup.Done()
			wp.work(ctx)
		}()
	}
}

func (wp *WorkerPool) Stop() {
	close(wp.exit)
	wp.waitGroup.Wait()
}

func (wp *WorkerPool) work(ctx context.Context) {
	defer func() {
		if err := recover(); err != nil {
			log.Error("recover:%v", err)
		}
	}()
	sub_ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	for {
		select {
		case <-wp.exit:
			return
		case <-ctx.Done():
			return
		case work := <-wp.workerCh:
			wp.dowork(sub_ctx, work)
		}
	}
}

func (wp *WorkerPool) dowork(ctx context.Context, work *Work) {
	//防止崩溃影响其他
	defer func() {
		if err := recover(); err != nil {
			log.Error("recover:%v", err)
		}
	}()
	work.worker.DoWork(ctx, work.args...)
}
